import math
from pathlib import Path
from typing import Any, Literal, Optional
from matplotlib.axes._axes import Axes
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MultipleLocator, ScalarFormatter
import xml.etree.ElementTree as ET

from .grouper import *
from .reader import *

PLOT_ALL_CHARGE = 0
PLOT_FAST_CHARGE = "快"
PLOT_SLOW_CHARGE = "慢"
PlotAccumMode = Literal[0,"快","慢"]
DEFAULT_SAVE_TO = ".\\results\\figures"
DEFAULT_CHN_FONT = 'SimSun' # 宋体

class NetPlot:
    @staticmethod
    def __set_style():
        from matplotlib import rcParams
        config = {
            "font.family":"serif",      # 有笔锋字体，没有笔锋是sans-serif
            "font.size":9,
            "mathtext.fontset":"stix",  # matplotlib渲染数学字体时使用的字体，和Times New Roman差别不大
            "font.serif":[DEFAULT_CHN_FONT],    
            "axes.unicode_minus":False  # 处理负号，即-号
        }
        rcParams.update(config)
        #sns.set_theme(style="white", font=DEFAULT_CHN_FONT, font_scale=1.4)

    @staticmethod
    def __pre(w,h,remove_edge=True)-> tuple[Figure, Axes]:
        fig, ax = plt.subplots(figsize=(w, h), dpi=300, constrained_layout=True)
        if remove_edge:
            ax.spines['right'].set_visible(False)  # 去除右边边框
            ax.spines['top'].set_visible(False)  # 去除上边边框
        return fig, ax
    
    @staticmethod
    def __post(fig,ax,xlabel,ylabel,title,save_to,mloc=True,zeroxmargin=True):
        if mloc:
            x_major_locator = MultipleLocator(60)
            ax.xaxis.set_major_locator(x_major_locator)
        if zeroxmargin: ax.margins(x=0)
        y_formatter = ScalarFormatter(useOffset=False)
        ax.yaxis.set_major_formatter(y_formatter)
        ax.set_xlabel(xlabel, fontweight="bold", font=DEFAULT_CHN_FONT)
        ax.tick_params(axis='x', labelrotation=270)
        ax.set_ylabel(ylabel, fontweight="bold", font=DEFAULT_CHN_FONT)
        ax.grid(False)
        ax.set_title(title, font=DEFAULT_CHN_FONT, fontsize=12)
        fig.savefig(save_to)
        plt.close()
    
    def __init__(self,data:ReadOnlyStatistics,veh_xml:Optional[str]=None,truncate:int=86400,save_to:str=DEFAULT_SAVE_TO):
        '''
        初始化
            data: ReadOnlyStatistics统计数据
            veh_xml: 车辆行程文件(可选)
            save_to: 图片默认保存路径
        '''
        self.__set_style()
        self.__save=save_to
        Path(save_to).mkdir(parents=True,exist_ok=True)
        # Statistics
        self.mySta=data
        self.grouper=MinAvgGrouper(data.TimeLine, truncate)
        # Trips
        self.__veh_xml=veh_xml
        self.__trip_dpt=None
    
    def __save_to(self,save_to):
        return save_to if save_to is not None else self.__save
    
    def __draw(self,ax,data=None):
        if data is not None: ax.plot(self.grouper.Process(data))
        cnt=len(self.grouper.MinuteList)
        k=max(1,60*((self.grouper.MinuteList[-1]-self.grouper.MinuteList[0])//(60*36)))
        ax.set_xticks(range(0,cnt,k))
        ax.set_xticklabels(self.grouper.TimeLine[::k])

    def cs_price(self,cs_name:str,save_to:Optional[str]=None,pbuy:bool=True,psell:bool=False):
        '''
        绘制充电站价格曲线
            cs_name: 充电站名称
            save_to: 保存路径, None则是NetPlot创建时定义的路径
        '''
        if FILE_CS not in self.mySta: raise ValueError("既有数据不支持绘制充电站价格")
        save_to=self.__save_to(save_to)
        fig, ax = self.__pre(12,3)
        if pbuy: ax.plot(self.grouper.Process(self.mySta.CS_pricebuy_of(cs_name)),label="用户购电价格")
        if psell: ax.plot(self.grouper.Process(self.mySta.CS_pricesell_of(cs_name)),label="用户售电价格")
        self.__draw(ax)
        fig.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax.transAxes)
        self.__post(fig,ax,"时间","价格/元",f"充电站{cs_name}价格",f"{save_to}\\充电站{cs_name}价格.png")
    
    def cs_prices(self,cs_names:Optional[list[str]]=None,save_to:Optional[str]=None,pbuy:bool=True,psell:bool=False,display:bool=True):
        '''
        绘制充电站购电售电价格曲线
            cs_name: 充电站名称
            save_to: 保存路径, None则是NetPlot创建时定义的路径
        '''
        save_to=self.__save_to(save_to)
        if cs_names is None: cs_names = self.mySta.CS_head
        for i,cs_name in enumerate(cs_names,start=1):
            if display : print(f"\r{i}/{len(cs_names)}: {cs_name}   ",end='')
            self.cs_price(cs_name,save_to,pbuy,psell if not cs_name.startswith("CS") else False)
        if display: print()

    def cs_curve(self,cs_name:str,save_to:Optional[str]=None,wcnt:bool=True,loadc:bool=True,loadd:bool=True,loadt:bool=True,loaddm:bool=True):
        '''
        绘制单个充电站(快充站/路边慢充)的曲线
            cs_name: 单个充电站名称
            save_to: 保存路径, None则是NetPlot创建时定义的路径
            wcnt: 是否绘制排队车辆数量曲线
            loadc: 是否绘制充电负荷曲线
            loadd: 是否绘制V2G实际放电负荷曲线
            loadt: 是否绘制净负荷曲线
            loaddm: 是否绘制V2G最大放电负荷曲线
        '''
        if FILE_CS not in self.mySta is None: raise ValueError("既有数据不支持绘制充电站曲线")
        save_to=self.__save_to(save_to)
        fig, ax1 = self.__pre(12,3)
        if wcnt and not loadc and not loadd and not loadt:
            self.__draw(ax1,self.mySta.CS_net_load_of(cs_name))
            yleft="车辆数"
        else:
            if wcnt:
                ax2 = ax1.twinx()
                assert isinstance(ax2, Axes)
                ax2.set_ylabel("车辆数", fontweight="bold", font=DEFAULT_CHN_FONT, fontsize=8)
                ax2.plot(self.grouper.Process(self.mySta.CS_count_of(cs_name)),label="站内车辆数",color="red",linestyle=":",linewidth=1)
            if loadt:
                ax1.plot(self.grouper.Process(self.mySta.CS_net_load_of(cs_name)),label="净充电负荷",color="aquamarine")
            if loadc:
                ax1.plot(self.grouper.Process(self.mySta.CS_charge_load_of(cs_name)),label="充电负荷",color="green",
                    linestyle="--" if loadt else "-",linewidth=1)
            if loadd:
                ax1.plot(self.grouper.Process(self.mySta.CS_v2g_load_of(cs_name)),label="V2G实际放电负荷",color="dodgerblue")
            if loaddm:
                ax1.plot(self.grouper.Process(self.mySta.CS_v2g_cap_of(cs_name)),label="V2G最大放电负荷",linestyle="-.",color="blue",linewidth=1)
            self.__draw(ax1)
            yleft="负荷/kW"
        fig.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax1.transAxes)
        self.__post(fig,ax1,"时间",yleft,f"充电站{cs_name}参数",f"{save_to}\\充电站{cs_name}参数.png")

    def cs_curves(self,cs_names:Optional[list[str]]=None,save_to:Optional[str]=None,display:bool=True,
                  wcnt:bool=True,loadc:bool=True,loadd:bool=True,loadt:bool=True,loaddm:bool=True):
        '''
        绘制所有充电站(快充站/路边慢充)的负荷曲线
            cs_names: 一个list[str], 为充电站名称列表, 如果为None, 则绘制所有; 
            save_to: 保存路径,  None则是NetPlot创建时定义的路径
            display: 是否显示绘制进度
        '''
        save_to=self.__save_to(save_to)
        if cs_names is None: cs_names=self.mySta.CS_head
        for i,cs_name in enumerate(cs_names,start=1):
            if display : print(f"\r{i}/{len(cs_names)}: {cs_name}   ",end='')
            if cs_name.startswith("CS"):
                self.cs_curve(cs_name,save_to,wcnt,loadc,False,False,False)
            else:
                self.cs_curve(cs_name,save_to,wcnt,loadc,loadd,loadt,loaddm)
        if display: print()

    def cs_load_accum(self,mode:PlotAccumMode=PLOT_ALL_CHARGE,save_to:Optional[str]=None):
        '''
        绘制快慢充面积图
            mode: PLOT_FAST_CHARGE为快充, PLOT_SLOW_CHARGE为慢充, PLOT_ALL_CHARGE则快慢充都绘制
            save_to: 保存路径, None则是NetPlot创建时定义的路径
        '''
        if FILE_CS not in self.mySta: raise ValueError("既有数据不支持绘制快慢充面积图")
        self.mySta.ForceLoad(FILE_CS)
        if mode==PLOT_FAST_CHARGE:
            cs_names = [x for x in self.mySta.CS_head if x.startswith("CS")]
        elif mode==PLOT_SLOW_CHARGE:
            cs_names = [x for x in self.mySta.CS_head if not x.startswith("CS")]
        elif mode==PLOT_ALL_CHARGE:
            self.cs_load_accum(PLOT_FAST_CHARGE,save_to)
            self.cs_load_accum(PLOT_SLOW_CHARGE,save_to)
            return 
        save_to=self.__save_to(save_to)
        x = self.grouper.TimeLine
        y = np.vstack([self.grouper.Process(self.mySta.CS_net_load_of(cs)) for cs in cs_names])
        fig, ax = self.__pre(15,4)
        stacks = ax.stackplot(x, y, labels=cs_names)
        self.__draw(ax)
        if mode==PLOT_FAST_CHARGE: ax.legend(loc="upper right")
        self.__post(fig,ax,"时间","负荷/kW",f"{mode}充负荷面积图",f"{save_to}\\{mode}充负荷面积图.png")

    def ev_attrib(self,veh_id:str,attrib:Union[str,list[str]],save_to:Optional[str]=None):
        f'''
        绘制电动汽车参数曲线
            veh_id: 车辆ID
            attrib: 参数名, 可选{EV_ATTRIB}
            save_to: 保存路径, None则是NetPlot创建时定义的路径
        '''
        if FILE_EV not in self.mySta: raise ValueError("既有数据不支持绘制EV曲线")
        save_to=self.__save_to(save_to)
        fig, ax = self.__pre(12,3)
        if isinstance(attrib, str): attrib = [attrib]
        if "cpure" in attrib or "earn" in attrib or "cost" in attrib:
            ax2 = ax.twinx()
            assert isinstance(ax2, Axes)
            ax2.set_ylabel("价格", fontweight="bold", font=DEFAULT_CHN_FONT, fontsize=8)
        ax.set_ylim(0,1)
        names = {
            "cpure": ("净成本","blue"),
            "earn": ("收入","red"),
            "cost": ("成本","green"),
            "soc": ("SoC","black"),
            "status": ("车辆状态","magenta"),
        }
        for attr in attrib:
            if attr == "cpure":
                data = self.mySta.EV_attrib_of(veh_id,"cost") - self.mySta.EV_attrib_of(veh_id,"earn")
            elif attr == "status":
                data = self.mySta.EV_attrib_of(veh_id,attr)/4
            else:
                data = self.mySta.EV_attrib_of(veh_id,attr)
            if attr in ["cpure","earn","cost"]:
                ax2.plot(self.grouper.Process(data),label=names[attr][0],color=names[attr][1])
            else:
                ax.plot(self.grouper.Process(data),label=names[attr][0],color=names[attr][1])
        self.__draw(ax)
        fig.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax.transAxes)
        self.__post(fig,ax,"时间","SoC",f"{veh_id}参数曲线",f"{save_to}\\电动汽车{veh_id}参数曲线.png")
    
    def gen_compare(self,save_to:Optional[str]=None):
        if not self.mySta.has_GEN(): raise ValueError("既有数据不支持绘制发电机成本")
        save_to = self.__save_to(save_to)
        fig, ax = self.__pre(12,3)
        gen_names=self.mySta.gen_head
        for gen in gen_names:
            ax.plot(self.grouper.Process(self.mySta.G_attrib_of(gen,"P")),label=f"{gen}有功",linewidth=1)
        self.__draw(ax)
        fig.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax.transAxes)
        self.__post(fig,ax,"时间","有功功率/MW",f"发电机有功比较",f"{save_to}\\发电机有功比较.png")

    def gen_total(self,save_to:Optional[str]=None):
        if not self.mySta.has_GEN(): raise ValueError("既有数据不支持绘制发电机成本")
        cost_data = self.mySta.G_total("totC")
        save_to=self.__save_to(save_to)
        fig, ax = self.__pre(12,3)
        ax2 = ax.twinx()
        assert isinstance(ax2, Axes)
        ax2.set_ylabel("成本/元/h", fontweight="bold", font=DEFAULT_CHN_FONT, fontsize=8)
        ax2.plot(self.grouper.Process(cost_data),label="总发电成本",color="red",linewidth=1)
        ax.plot(self.grouper.Process(self.mySta.G_total("totP")),label="总发电有功功率",color="green",linewidth=1)
        ax.plot(self.grouper.Process(self.mySta.G_total("totQ")),label="总发电无功功率",linestyle="--",color="aquamarine",linewidth=1)
        self.__draw(ax)
        fig.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax.transAxes)
        self.__post(fig,ax,"时间","发电功率/MVA",f"单位时间发电总参数",f"{save_to}\\单位时间发电总参数.png")

        tl = self.grouper.SecondList
        lastT = tl[0] - (tl[1] - tl[0])
        total_cost = 0
        for t,d in zip(tl,cost_data[self.grouper.TruncatePos:]):
            total_cost += d * (t - lastT) / 3600
            lastT = t
        print(f"发电总成本: {total_cost:.2f}元")
        with open(f"{save_to}\\发电总成本{total_cost:.2f}元.txt","w",encoding="utf-8") as f:
            f.write(f"发电总成本: {total_cost:.2f}元")

    def gen_curve(self,gen_id:str,save_to:Optional[str]=None):
        f'''
        绘制发电机参数曲线
            gen_id: 发电机ID
            attrib: 参数名, 可选{GEN_ATTRIB}
            save_to: 保存路径, None则是NetPlot创建时定义的路径
        '''
        if not self.mySta.has_GEN(): raise ValueError("既有数据不支持绘制发电机曲线")
        save_to = self.__save_to(save_to)
        fig, ax1 = self.__pre(12,3)
        ax2 = ax1.twinx()
        ax2.set_ylabel("发电价格(元/h)", fontweight="bold", font=DEFAULT_CHN_FONT, fontsize=8)
        assert isinstance(ax2, Axes)
        ax2.plot(self.grouper.Process(self.mySta.G_attrib_of(gen_id,"costp")),label="发电价格",color="red",linewidth=1)
        ax1.plot(self.grouper.Process(self.mySta.G_attrib_of(gen_id,"P")),label="有功功率",color="blue")
        ax1.plot(self.grouper.Process(self.mySta.G_attrib_of(gen_id,"Q")),label="无功功率",color="green",linewidth=1)
        self.__draw(ax1)
        yleft="发电量/MVA"
        fig.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax1.transAxes)
        self.__post(fig,ax1,"时间",yleft,f"发电机{gen_id}参数",f"{save_to}\\发电机{gen_id}参数.png")

    def gen_curves(self,gen_names:Optional[list[str]]=None,save_to:Optional[str]=None,display:bool=True):
        '''
        绘制所有发电机的参数曲线
            gen_names: 一个list[str], 为充电站名称列表, 如果为None, 则绘制所有; 
            save_to: 保存路径,  None则是NetPlot创建时定义的路径
            display: 是否显示绘制进度
        '''
        save_to=self.__save_to(save_to)
        if gen_names is None: gen_names=self.mySta.gen_head
        if FILE_GEN not in self.mySta: raise ValueError("既有数据不支持绘制发电机曲线")
        for i, gen_name in enumerate(gen_names,start=1):
            if display : print(f"\r{i}/{len(gen_names)}: {gen_name}   ",end='')
            self.gen_curve(gen_name,save_to)
        if display: print()

    def __load_trips(self,veh_xml_file:Optional[str]):
        if self.__trip_dpt is not None: return
        if veh_xml_file is None: veh_xml_file=self.__veh_xml
        if veh_xml_file is None: raise ValueError("未指定trips文件名")
        root = ET.ElementTree(file=veh_xml_file).getroot()
        if root is None: raise ValueError("行程文件未指定或已损坏")
        self.__trip_dpt = {}
        self.__trip_maxlen = 0
        for vehicle in root:
            val=[trip.attrib["depart"] for trip in vehicle]
            self.__trip_dpt[vehicle.attrib["id"]]=val
            self.__trip_maxlen = max(self.__trip_maxlen,len(val))
    
    def bus_curve(self,bus_id:str,save_to:Optional[str]=None):
        if FILE_BUS not in self.mySta: raise ValueError("既有数据不支持绘制母线曲线")
        save_to=self.__save_to(save_to)
        fig, ax1 = self.__pre(12,3)
        if bus_id+"#cshd" in self.mySta.Headof(FILE_BUS):
            ax2 = ax1.twinx()
            ax2.set_ylabel("电价(元/kWh)", fontweight="bold", font=DEFAULT_CHN_FONT, fontsize=8)
            assert isinstance(ax2, Axes)
            ax2.plot(self.grouper.Process(self.mySta.bus_attrib_of(bus_id,"cshd")),label="影子电价",color="red")
            ax1.plot(self.grouper.Process(self.mySta.bus_attrib_of(bus_id,"Pg")),label="有功发电",color="green",linewidth=1)
            ax1.plot(self.grouper.Process(self.mySta.bus_attrib_of(bus_id,"Qg")),label="无功发电",linestyle="--",color="aquamarine",linewidth=1)
        ax1.plot(self.grouper.Process(self.mySta.bus_attrib_of(bus_id,"Pd")),label="有功负荷",color="blue",linewidth=1)
        ax1.plot(self.grouper.Process(self.mySta.bus_attrib_of(bus_id,"Qd")),label="无功负荷",linestyle="--",color="dodgerblue",linewidth=1)
        self.__draw(ax1)
        yleft="发电量/MVA"
        fig.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax1.transAxes)
        self.__post(fig,ax1,"时间",yleft,f"母线{bus_id}参数",f"{save_to}\\母线{bus_id}参数.png")

    def bus_curves(self,bus_names:Optional[list[str]]=None,save_to:Optional[str]=None,display:bool=True):
        '''
        绘制所有母线的参数曲线
            gen_names: 一个list[str], 为充电站名称列表, 如果为None, 则绘制所有; 
            save_to: 保存路径,  None则是NetPlot创建时定义的路径
            display: 是否显示绘制进度
        '''
        save_to=self.__save_to(save_to)
        if bus_names is None: bus_names=self.mySta.bus_head
        if FILE_GEN not in self.mySta: raise ValueError("既有数据不支持绘制母线曲线")
        for i,bus_name in enumerate(bus_names,start=1):
            if display : print(f"\r{i}/{len(bus_names)}: {bus_name}   ",end='')
            self.bus_curve(bus_name,save_to)
        if display: print()

    def bus_total(self,save_to:Optional[str]=None):
        if not self.mySta.has_BUS(): raise ValueError("既有数据不支持绘制母线总参数")
        save_to=self.__save_to(save_to)
        fig, ax = self.__pre(12,3)
        ax.plot(self.grouper.Process(self.mySta.bus_total("totPd")),label="总有功负荷",color="blue",linewidth=1)
        ax.plot(self.grouper.Process(self.mySta.bus_total("totQd")),label="总无功负荷",linestyle="--",color="dodgerblue",linewidth=1)
        ax.plot(self.grouper.Process(self.mySta.bus_total("totPg")),label="总有功发电",color="green",linewidth=1)
        ax.plot(self.grouper.Process(self.mySta.bus_total("totQg")),label="总无功发电",linestyle="--",color="aquamarine",linewidth=1)
        self.__draw(ax)
        ax.legend(loc="upper right", bbox_to_anchor=(1, 1), bbox_transform=ax.transAxes)
        self.__post(fig,ax,"时间","功率/MVA",f"单位时间母线总参数",f"{save_to}\\单位时间母线总参数.png")

    def trip_hist(self,trip_nums:Optional[list[int]]=None,veh_xml_file:Optional[str]=None,save_to:Optional[str]=None):
        '''
        统计第i次行程的车辆出发时间概率分布
            trip_nums: 一个列表, 列表中每个值表示第i次行程, 从1开始: 1, 2, 3, ..., None表示所有行程
            veh_xml_file: 车辆行程文件, 如果定义NetPlot时已提供, 则可以省略
            save_to: 保存路径,  None则是NetPlot创建时定义的路径
        '''
        self.__load_trips(veh_xml_file)
        if self.__trip_dpt is None: raise ValueError("Trips未加载")
        save_to=self.__save_to(save_to)
        if trip_nums is None: trip_nums=[i for i in range(1,self.__trip_maxlen+1)]
        for trip_num in trip_nums:
            fig, ax = self.__pre(12,3)
            triptime = [float(self.__trip_dpt[x][trip_num-1]) for x in self.__trip_dpt.keys() if len(self.__trip_dpt[x])>trip_num-1]
            minh=math.floor(min(triptime)/3600);maxh=math.ceil(max(triptime)/3600)
            bins=range(minh*3600, maxh*3600+1,3600)
            ax.hist([triptime], bins=bins, color='#5799C6', edgecolor="black", alpha=0.8)  # color='#5799C6'
            ax.set_xticks(bins)
            ax.set_xticklabels([f"{x:02}:00" for x in range(minh, maxh+1)])
            self.__post(fig,ax,"时间","出发车辆数目",f"第{trip_num}次行程不同时间段出发车辆数目",f"{save_to}\\plot_trip_{trip_num}.png",False)
        
    def ev_route(self,veh_id:str,save_to:Optional[str]=None):
        '''
        绘制路径
            veh_id: 车辆ID
            route_xml_file: 路径xml文件, 如果定义NetPlot时已提供, 则可以省略
            save_to: 保存路径,  None则是NetPlot创建时定义的路径
        '''
        if FILE_EV not in self.mySta: raise ValueError("既有数据不支持绘制EV路径")
        save_to=self.__save_to(save_to)
        # 将数据点按顺序绘制
        x_values, y_values = self.mySta.EV_loc_of(veh_id)
        fig, ax = self.__pre(8,8,False)
        plt.plot(x_values, y_values, linestyle='-')

        self.__post(fig,ax,"X","Y",f"{veh_id}的行程轨迹",f"{save_to}\\{veh_id}的行程轨迹.png",False, False)
